/* Copyright (C) 2005, 2006, 2007, 2008, 2009, 2010 Mario Orsi
   This file is part of Brahms.
   Brahms is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 3 of the License, or (at your option) any later version.
   Brahms is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
   You should have received a copy of the GNU General Public License along with Brahms.  If not, see <http://www.gnu.org/licenses/>. */

// forceBdMain.c -- various functions that evaluate bonded (intramolecular) interactions
// Warning: many functions and variables are specific to the lipid topology - see dataStructs.h

//#define DEBUG // uncomment this line to turn on debugging functionalities

#include "dataStructs.h"
#include "definitions.h"

/* External data references: */
extern const VecR *siteOrientVec; 
extern const VecR *siteCoords;
extern VecR *siteForce;
extern VecR *siteTorque;
extern const VecR region;
extern Site *site;  
extern const int lpp, nSites, stepCount, stepLpp; 
extern const real springRigidity, **sig;

// External function declarations:
void DistributeVirial( RMat, int, int ); 
void PrintMatrix( RMat ); 

/* Global variables: */
real scaleRefSpringLength; // input parameter: referenceSpringLength_ij = scaleRefSpringLength * (sigma_i+sigma_j)/2
real refAngle_PhosGlyEst; /* input parameter: reference angle for the phosphate-glycerol-ester triplet */
real refAngle_CholPhosGly; /* input parameter: reference angle for phosphocholine headgroup tilt */
real angleRigidity; /* input parameter: rigidity of the angle potential for all angles */ 
real refAngle_cisUnsat; /* reference angle for tail sites corresponding to triplets involving a cis unsaturated double bond */
real orientRigidity; /* input parameter: orientational rigidity for dipolar lipid sites */ 

/* CalcAngleBendInteraction -- calculates angle bending interaction, adapted from [Rapaport, pp.283-285] */
static void CalcSingleAngleBendInteraction( real *potentialEnergy,     /* out - harmonic potential */
					    RMat *virial,              /* out - virial */
					    VecR *force1,              /* out - force on site "1", preceding the angle */
					    VecR *force2,              /* out - force on site "2", after the angle  */
					    const real rigidity,       /* in - force constant of the angle potential */
					    const real referenceAngle, /* in - reference (or equilibrium) angle */
					    const int i1,              /* in - index of site "1", preceding the angle */
					    const int iA,              /* in - index of site "A", at the angle */
					    const int i2 )             /* in - index of site "2", after the angle */
{
  VecR rA1, /* rA - r1, distance vector from site "1" to site "A" */
    r2A,    /* r2 - rA, distance vector from site "A" to site "2" */
    r1A;    /* r1 - rA, distance vector from site "A" to site "2" */
  real c11, c12, c22, cd, cos_PiMinusReferenceAngle, cos_PiMinusActualAngle, f;
  RMat virial_1A, virial_2A;  

  cos_PiMinusReferenceAngle = cos( PI - referenceAngle );
  VSub( rA1, siteCoords[ iA ], siteCoords[ i1 ] ); 
  VWrapAll( rA1 ); /* minimum image convention */
  VSub( r2A, siteCoords[ i2 ], siteCoords[ iA ] ); 
  VWrapAll( r2A ); /* minimum image convention */
  c11 = VLenSq( rA1 );
  c12 = VDot( rA1, r2A );
  c22 = VLenSq( r2A );
  cd = sqrt( c11 * c22 );
  cos_PiMinusActualAngle = c12 / cd;
  f = - rigidity * ( cos_PiMinusActualAngle - cos_PiMinusReferenceAngle );
  VSSAdd( *force1, c12 / c11, rA1, -1., r2A );
  VScale( *force1, f / cd ); 
  VSSAdd( *force2, 1., rA1, - c12 / c22, r2A );
  VScale( *force2, f / cd );
  *potentialEnergy = 0.5 * rigidity * Sqr( cos_PiMinusActualAngle - cos_PiMinusReferenceAngle );

  /* virial contribution: */
  VSCopy( r1A, -1., rA1 ); 
  VOuter( virial_1A.u, r1A, *force1 );
  VOuter( virial_2A.u, r2A, *force2 );
  MCopy( virial->u, virial_1A.u );
  MMAdd( virial->u, virial_2A.u );
  if ( (lpp) && ( stepCount % stepLpp == 0 ) ) {
    DistributeVirial( virial_1A, i1, iA );
    DistributeVirial( virial_2A, i2, iA );
  }
  
/*   PrintMatrix( virial_1A ); */
/*   PrintMatrix( virial_2A ); */
/*   PrintMatrix( *virial ); */
/*   printf("\n"); */
  // checking that trace of the virial is zero [in 't Veld and Rutledge, Macromolecules 2003, 36, 7358]
/*   printf("### trace(W_1a) = % .50f\n",  Trace( virial_1A.u )); */
/*   printf("### trace(W_2a) = % .50f\n",  Trace( virial_2A.u )); */
/*   printf("### trace(Wtot) = % .50f\n\n",Trace( virial->u )); */
}

/***************************************************************************************************
 * ComputeAngleBendInteractions -- evaluates all angle bending interactions within lipid molecules *
 ***************************************************************************************************/
void ComputeAngleBendInteractions( real *potentialEnergySum, RMat *systemVirial )
{
  int J1, jAngle, J2, n;
  real potentialEnergy;
  RMat virial, virialSum;
  real refAngle, /* used to store the reference angle of a particular interaction */
    rigidity; /* used to store the rigidity constant of a particular interaction */
  VecR force1, force2;
  const real refAngle_sat = PI; /* reference angle for tail sites corresponding to triplets connected by saturated bonds */

  MZero( virialSum.u ); /* initialize virial accumulation */

  rigidity = angleRigidity; /* setting rigidity to a single value used for all angles (input parameter 'angleRigidity') */
  
  DO_SITE {
    switch ( site[ n ].lipidUnit ) { /* identify the site "n" corresponding to the angle "(n-1)--(n)--(n+1)" */
    case WATER: 
    case CHOLINE_OR_AMINE:
    case TAIL_A5: 
    case TAIL_B5: 
      break; /* there is no angle potential centered on these sites */
    case GLYCEROL: /* calculating angle bending for phosphate(2)-glycerol(3)-ester_B(10) -- SPECIAL CASE */
      /* setting indexes based on the DOPC topology as specified in dataStruct.h */  
      J1 = n-1; /* phosphate */
      jAngle = n; /* glycerol */
      J2 = n+7; /* ester_b */
      /* setting reference angle and rigidity */
      refAngle = refAngle_PhosGlyEst; 
      CalcSingleAngleBendInteraction( &potentialEnergy, &virial, &force1, &force2, rigidity, refAngle, J1, jAngle, J2 );
      *potentialEnergySum += potentialEnergy;
      MMAdd( virialSum.u, virial.u );
      VVAdd( siteForce[J1],     force1 );
      VVSub( siteForce[jAngle], force1 ); /* Newton 3rd law */
      VVSub( siteForce[jAngle], force2 ); /* Newton 3rd law */
      VVAdd( siteForce[J2],     force2 );
      /*       /\* now calculating angle bending for ester_A(4)-glycerol(3)-ester_B(10) -- SPECIAL CASE *\/ */
      /*       /\* setting indexes based on the DOPC topology as specified in dataStruct.h *\/   */
      /*       J1 = n+1; /\* resetting J1 index to current site, i.e., ester_A *\/  */
      /*       /\* setting reference angle and rigidity *\/ */
      /*       refAngle = refAngle_EstGlyEst;  */
      /*       CalcSingleAngleBendInteraction( &potentialEnergy, &virial, &force1, &force2, rigidity, refAngle, J1, jAngle, J2 ); */
      /*       *potentialEnergySum += potentialEnergy; */
      /*       MMAdd( virialSum.u, virial.u ); */
      /*       VVAdd( siteForce[J1],     force1 ); */
      /*       VVSub( siteForce[jAngle], force1 ); /\* Newton 3rd law *\/ */
      /*       VVSub( siteForce[jAngle], force2 ); /\* Newton 3rd law *\/ */
      /*       VVAdd( siteForce[J2],     force2 ); */
      /* fall through to the general case so that also phosphate(2)-glycerol(3)-ester_A(4) will be evaluated */
    case PHOSPHATE: /* choline(1)-phosphate(2)-glycerol(3) */
    case ESTER_A:   /* glycerol(3)-ester_A(4)-hydrocarbon_a1(5) */
    case ESTER_B:   /* glycerol(3)-ester_B(10)-hydrocarbon_b1(11) -- SPECIAL CASE */
    case TAIL_A1:   /* ester_A(4)-hydrocarbon_a1(5)-hydrocarbon_a2(6) */
    case TAIL_A2:   /* hydrocarbon_a1(5)-hydrocarbon_a2(6)-hydrocarbon_a3(7) */
    case TAIL_A3:   /* hydrocarbon_a2(6)-hydrocarbon_a3(7)-hydrocarbon_a4(8) */
    case TAIL_A4:   /* hydrocarbon_a3(7)-hydrocarbon_a4(8)-hydrocarbon_a5(9) */
    case TAIL_B1:   /* ester_B(4)-hydrocarbon_b1(5)-hydrocarbon_b2(6) */
    case TAIL_B2:   /* hydrocarbon_b1(5)-hydrocarbon_b2(6)-hydrocarbon_b3(7) */
    case TAIL_B3:   /* hydrocarbon_b2(6)-hydrocarbon_b3(7)-hydrocarbon_b4(8) */
    case TAIL_B4:   /* hydrocarbon_b3(7)-hydrocarbon_b4(8)-hydrocarbon_b5(9) */
      /* setting indexes based on the DOPC topology as specified in dataStruct.h */
      jAngle = n;
      J2 = n+1;
      if ( site[ n ].lipidUnit == ESTER_B ) /* special case for glycerol-ester_b-hydrocarbon_b1 */
	J1 = n-7; /* 7 is the number of units separating ester_b from glycerol - see dataStruct.h */
      else J1 = n-1; /* general case */
      switch ( site[ n ].lipidUnit ) { /* setting reference angles and rigidity constants */
      case PHOSPHATE:
	refAngle = refAngle_CholPhosGly;
	break;
      case GLYCEROL: /* phosphate-glycerol-esterA */
	refAngle = refAngle_PhosGlyEst; 
	break;
      case TAIL_A2:
      case TAIL_B2:
	refAngle = refAngle_cisUnsat; /* set reference angle that models cis unsaturation */
	break;
      default:/* general case */
	refAngle = refAngle_sat;
	break;
      }
      CalcSingleAngleBendInteraction( &potentialEnergy, &virial, &force1, &force2, rigidity, refAngle, J1, jAngle, J2 );
      *potentialEnergySum += potentialEnergy;
      MMAdd( virialSum.u, virial.u );
      VVAdd( siteForce[J1],     force1 );
      VVSub( siteForce[jAngle], force1 ); /* Newton 3rd law */
      VVSub( siteForce[jAngle], force2 ); /* Newton 3rd law */
      VVAdd( siteForce[J2],     force2 );
      break;
    default:
      printf("Error: default case reached in CalcAngleBendInteractions.\n"); exit(0);
    }
  }
  MMAdd( systemVirial->u, virialSum.u ); /* add virial from angle bending (virialSum) to total system virial */
}

static void CalcOrientRestrainInteraction( real *potEn, // out - potential energy 
					   VecR *torque, // out - torque
					   const real k, // in - rigidity constant
					   const VecR eRef, // in - reference orientation
					   const VecR eAct ) // in - actual orientation
{
  real cosTheta = VDot( eRef, eAct );

  *potEn = 0.5 * k * Sqr( cosTheta - 1 );
  VCross( *torque, eAct, eRef );
  VScale( *torque, k * ( 1 - cosTheta ) );
}

void ComputeOrientationRestraints( real *potEnSum )
{
  int iPhos, iGly, iEst, iHydro, n;
  real potEn;
  VecR torque;
  VecR rPhosGly, rEstHydro, rVersor; 

  DO_SITE {
    if ( site[ n ].lipidUnit == GLYCEROL ) { // restrain glycerol orientation
      iGly = n;
      iPhos = n - 1;
      VSub( rPhosGly, siteCoords[ iPhos ], siteCoords[ iGly ] ); // rGly->rPhos vector
      VWrapAll( rPhosGly ); /* Minimum Image Convention */
      VSCopy( rVersor, 1 / VLen( rPhosGly ), rPhosGly ); // gly->phos unit vector (rPhosGly = rPhos - rGly)
      CalcOrientRestrainInteraction( &potEn, &torque, orientRigidity, rVersor, siteOrientVec[iGly] );
      *potEnSum += potEn;
      VVAdd( siteTorque[ iGly ], torque ); /* add to global array */
    } else if ( (site[n].lipidUnit == ESTER_A) || (site[n].lipidUnit == ESTER_B) ) { 
      // restrain orientation of ester sites to hydro_1 -> ester vector:
      iEst = n;
      iHydro = n + 1;
      VSub( rEstHydro, siteCoords[ iEst ], siteCoords[ iHydro ] );
      VWrapAll( rEstHydro ); /* Minimum Image Convention */
      VSCopy( rVersor, 1 / VLen( rEstHydro ), rEstHydro ); // unit vector along the tail->est direction
      CalcOrientRestrainInteraction( &potEn, &torque, orientRigidity, rVersor, siteOrientVec[iEst] );
      *potEnSum += potEn;
      VVAdd( siteTorque[ iEst ], torque ); // add to global array
    }
  }
}

static void CalcHarmonicInteraction( real *potentialEnergy, /* out - harmonic potential */
				     RMat *virial, /* out - virial */
				     VecR *harmonicForce12, /* out - force */
				     const real kappaSpring, /* in - rigidity constant */
				     real springReferenceLength, /* in - reference (or equilibrium) length */ 
				     const VecR rSpring12, /* "spring vector" r1-r2 (wrapped as per MIC)  */
				     const int i1, /* index of site1 in the global site array */
				     const int i2 ) /* index of site2 in the global site array */
{  
  real rSpring12Length, deltaSpringLength; 
  VecR springVersor; /* unit spring vector */

  springReferenceLength *= scaleRefSpringLength; /* scaling according to input parameter 'scaleRefSpringLength' */
  
  rSpring12Length = VLen( rSpring12 );
  deltaSpringLength = rSpring12Length - springReferenceLength; 
  *potentialEnergy = 0.5 * kappaSpring * Sqr( deltaSpringLength );
  VSCopy( springVersor, 1./rSpring12Length, rSpring12 );
  VSCopy( *harmonicForce12, -kappaSpring * deltaSpringLength, springVersor );

  /* Virial contribution: */
  VOuter( virial->u, rSpring12, *harmonicForce12 ); // assuming springs are anchored at the sites' centres
  if ( (lpp) && ( stepCount % stepLpp == 0 ) ) DistributeVirial( *virial, i1, i2 );
}

void ComputeBondStretchInteractions( real *potEnSum, RMat *systemVirial ) 
{
  int i1, i2, iGly, iEst, n;
  real potEn, springReferenceLength;
  RMat virial, virialSum;
  VecR force12;
  VecR r12, forceEstGly, rEstGly; 

  MZero( virialSum.u );

  DO_SITE {
    switch ( site[ n ].lipidUnit ) {
    case WATER:     /* break - nonlipid site */ 
    case TAIL_A5: // break - this case is considered in TAIL_A4
    case TAIL_B5: // break - this case is considered in TAIL_B4
      break;
    case ESTER_B: // ESTER_B***TAIL_B1 special case
      iEst = n; 
      iGly = n - 7; 
      springReferenceLength = sig[site[iGly].type][site[iEst].type]; /* setting reference length to (sigma1+sigma2)/2 */
      VSub( rEstGly, siteCoords[ iEst ], siteCoords[ iGly ] ); // glycerol->ester COM interdistance vector
      VWrapAll( rEstGly ); /* Minimum Image Convention */
      /* calculate harmonic interaction between ester and glycerol mass centres */
      CalcHarmonicInteraction( &potEn, &virial, &forceEstGly, springRigidity, springReferenceLength, rEstGly, iEst, iGly );
      *potEnSum += potEn;
      MMAdd( virialSum.u, virial.u );
      VVAdd( siteForce[ iEst ], forceEstGly ); // adding force on ester due to glycerol
      VVSub( siteForce[ iGly ], forceEstGly );  /* Newton's Third Law... */
      // fall through to consider also ESTER_B***TAIL_B1
    case CHOLINE_OR_AMINE: /* choline***phosphate (DOPC) or amine***phosphate (DOPE) */
    case PHOSPHATE: /* phosphate***glycerol */
    case GLYCEROL: /* GLYCEROL***ESTER_A */  
    case ESTER_A: // ESTER_A***TAIL_A1

    case TAIL_A1: // TAIL_A1***TAIL_A2
    case TAIL_B1: // TAIL_B1***TAIL_B2

    case TAIL_A2: // TAIL_A2***TAIL_A3
    case TAIL_B2: // TAIL_B2***TAIL_B3

    case TAIL_A3: // TAIL_A3***TAIL_A4
    case TAIL_B3: // TAIL_B3***TAIL_B4

    case TAIL_A4: // do TAIL_A4***TAIL_A5
    case TAIL_B4: // do TAIL_B4***TAIL_B5
      /* consider here all harmonic pair interactions *between mass centres* for which i1=n and i2+n+1 */
      i1 = n;
      i2 = n + 1;
      springReferenceLength = sig[site[i1].type][site[i2].type]; /* setting reference length to (sigma1+sigma2)/2 */
      VSub( r12, siteCoords[ i1 ], siteCoords[ i2 ] ); // r2->r1 vector
      VWrapAll( r12 ); /* Minimum Image Convention */
      CalcHarmonicInteraction( &potEn, &virial, &force12, springRigidity, springReferenceLength, r12, i1, i2 );
      *potEnSum += potEn;
      MMAdd( virialSum.u, virial.u );
      VVAdd( siteForce[ i1 ], force12 );
      VVSub( siteForce[ i2 ], force12 );  /* Newton's Third Law... */
      break;
    default:
      printf("Error: default case reached in ComputeBondStretchInteractions. Exiting...\n");
      exit(0);
    }    
  }
  MMAdd( systemVirial->u, virialSum.u ); /* add virial from bond stretching (virialSum) to total system virial */
}
